package com.kanq.documentLoader;

import com.kanq.common.utils.HanlpUtils;

import java.util.ArrayList;
import java.util.List;

public class SemanticChunker {

    /**
     * 改进的语义分块算法，结合句子边界检测和动态窗口调整
     *
     * @param text         需要分块的文本
     * @param maxChunkSize 最大分块大小
     * @param minChunkSize 最小分块大小
     * @return 分块后的文本列表
     */
    public static List<String> semanticChunking(String text, int maxChunkSize, int minChunkSize) {
        List<String> chunks = new ArrayList<>();

        if (text == null || text.isEmpty()) {
            return chunks;
        }

        if (maxChunkSize <= 0 || minChunkSize <= 0 || maxChunkSize < minChunkSize) {
            throw new IllegalArgumentException("分块大小参数不合法");
        }

        // 1. 使用HanLP进行句子级分割
        List<String> sentences = HanlpUtils.extractSentences(text);

        // 2. 合并短句子，分割长句子
        StringBuilder currentChunk = new StringBuilder();
        for (String sentence : sentences) {
            // 如果当前句子过长，单独处理
            if (sentence.length() > maxChunkSize) {
                // 先保存已积累的内容
                if (currentChunk.length() > 0) {
                    chunks.add(currentChunk.toString());
                    currentChunk.setLength(0);
                }
                // 对长句子进行滑动窗口分割
                chunks.addAll(splitLongSentence(sentence, maxChunkSize, minChunkSize));
            }
            // 如果加入当前句子不会超过最大大小
            else if (currentChunk.length() + sentence.length() <= maxChunkSize) {
                if (currentChunk.length() > 0) {
                    currentChunk.append(" ");
                }
                currentChunk.append(sentence);
            }
            // 如果加入当前句子会超过最大大小，但当前积累已经达到最小大小
            else if (currentChunk.length() >= minChunkSize) {
                chunks.add(currentChunk.toString());
                currentChunk.setLength(0);
                currentChunk.append(sentence);
            }
            // 否则继续积累
            else {
                if (currentChunk.length() > 0) {
                    currentChunk.append(" ");
                }
                currentChunk.append(sentence);
            }
        }

        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(currentChunk.toString());
        }

        return chunks;
    }

    /**
     * 对长句子进行滑动窗口分割
     */
    private static List<String> splitLongSentence(String sentence, int maxSize, int minSize) {
        List<String> subChunks = new ArrayList<>();
        int start = 0;
        int sentenceLength = sentence.length();

        while (start < sentenceLength) {
            // 尝试找到最近的句子边界（标点符号或空格）
            int end = findNearestBoundary(sentence, start, maxSize);

            // 如果找不到合适的边界，或者剩余部分小于最小大小，直接取到末尾
            if (end <= start || (sentenceLength - start) < minSize) {
                end = sentenceLength;
            }

            subChunks.add(sentence.substring(start, end).trim());
            start = end;
        }

        return subChunks;
    }

    /**
     * 在指定范围内查找最近的句子边界
     */
    private static int findNearestBoundary(String text, int start, int maxDistance) {
        int end = Math.min(start + maxDistance, text.length());

        // 优先查找句子结束标点
        for (int i = end - 1; i > start; i--) {
            char c = text.charAt(i);
            if (c == '.' || c == '。' || c == '!' || c == '！' || c == '?' || c == '？') {
                return i + 1;
            }
        }

        // 其次查找逗号或分号
        for (int i = end - 1; i > start; i--) {
            char c = text.charAt(i);
            if (c == ',' || c == '，' || c == ';' || c == '；') {
                return i + 1;
            }
        }

        // 最后查找空格
        for (int i = end - 1; i > start; i--) {
            if (Character.isWhitespace(text.charAt(i))) {
                return i + 1;
            }
        }

        return end; // 没找到边界，返回最大距离
    }
}